from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

import pickle as pkl
import pandas as pd
import random
import numpy as np
import time

import datetime
import math
import numpy as np
import torch
from torch import nn
from torch.nn import Module, Parameter
import torch.nn.functional as F

with open('./Data/Processed/word_emb_tensor.pkl', 'rb') as f:
    word_embed_tensor = pkl.load(f)
word_embed_tensor = torch.tensor(word_embed_tensor, device='cuda:4')

topk_index = []
n = word_embed_tensor.shape[0]
cos_sim = nn.CosineSimilarity(dim=1, eps=1e-6)
for i in range(n):
    if i%1000==0: print(i)
    dist0 = cos_sim(word_embed_tensor[i:i+1], word_embed_tensor[0:n//2])
    dist1 = cos_sim(word_embed_tensor[i:i+1], word_embed_tensor[n//2:])
    dist = torch.cat([dist0, dist1])
    ind = torch.topk(torch.tensor(dist), 100)[1]
    topk_index.append(ind)
topk_index = torch.stack(topk_index)
print(topk_index.shape)

topk_index_ = topk_index.cpu().numpy()

with open('./Data/Processed/topk_index_100.pkl', 'wb') as f:
    pkl.dump(topk_index_, f, pkl.HIGHEST_PROTOCOL)